package edu.gatech.cc.liam.marl.decpomdp;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map.Entry;
import java.util.Set;
import java.util.TreeSet;

import edu.gatech.cc.liam.core.Globals;
import edu.gatech.cc.liam.core.PairComparable;
import edu.gatech.cc.liam.core.rl.models.CartesianIterator;
import edu.gatech.cc.liam.core.rl.models.JointIterator;
import edu.gatech.cc.liam.core.rl.models.decpomdp.DecMDPSingleStagePolicy;
import edu.gatech.cc.liam.geometry.linear.NPoint;
import edu.gatech.cc.liam.geometry.linear.SparsePoint;

public class Belief {
	private static final long serialVersionUID = -3801572626724038881L;
	private SparsePoint dist; 
	private double valueAtBelief = Double.NaN;
	private ValueVector supportVector = null;
	public boolean hasImproved;
	
	class DVO extends PairComparable<Double, ValueVector> {
		private DecMDPSingleStagePolicy theActions;
		public DVO(double d, ValueVector vv) {
			super(d, vv);
			theActions = Belief.this.supportVector.nextActions;
		}
		public Double getHeuristic() { return this.obj1; }
		public ValueVector getSuccVV() { return this.obj2; }
		public DecMDPSingleStagePolicy getActions() { return theActions; }
		@Override
		public int hashCode() {
			return System.identityHashCode(this);
		}
		@Override
		public boolean equals(Object obj) {
			return this == obj;
		}
//		public double[] getObj() { return this.obj2.obj2; }
	}
//	private ArrayList<DVO> succValues;
	private HashMap<ValueVector, DVO> succValues = new HashMap<ValueVector, Belief.DVO>();
	private TreeSet<DVO> orderedSuccValues = new TreeSet<DVO>(Collections.reverseOrder());
	
	public DVO getBestSucc() {
		return orderedSuccValues.first();
	}
	public DVO removeBestSucc() {
		return orderedSuccValues.pollFirst();
	}
	public double getSuccImprovement() {
		return getBestSucc().obj1 - getValue();
	}
	public void updateSuccValues(Iterable<ValueVector> vvs) {
		HashMap<ValueVector, DVO> oldSuccVals = succValues;
		succValues = new HashMap<ValueVector, Belief.DVO>();
		orderedSuccValues.clear();
		for(ValueVector vv : vvs) {
			DVO e = oldSuccVals.get(vv);
			if(e != null) {
				succValues.put(vv, e);
				orderedSuccValues.add(e);
			} else {
				addSuccValue(vv);
			}
		}
	}
	public void addSuccValue(ValueVector vv) {
		double val = vv.getValueIfSuccessor(this);
		DVO e = new DVO(val, vv);
		succValues.put(vv, e);
		orderedSuccValues.add(e);
	}
	public void replaceSuccIfBetter(ValueVector vv) {
		double val = vv.getValueIfSuccessor(this);
		if(val > getBestSucc().obj1) {
			DVO e = new DVO(val, vv);
			succValues.clear();
			succValues.put(vv, e);
			orderedSuccValues.clear();
			orderedSuccValues.add(e);
		}
	}
	
	public void updateSuccValue(ValueVector oldVV, ValueVector newVV) {
		orderedSuccValues.remove(succValues.get(oldVV));
		succValues.remove(oldVV);
		double val = newVV.getValueIfSuccessor(this);
		DVO e = new DVO(val, newVV);
		succValues.put(newVV, e);
		orderedSuccValues.add(e);
	}
	public ValueVector getSuppFromBestSucc(DMDPValueFunction vf) { // if action isn't changed, find best successor policy
		DVO bestSucc = getBestSucc();
		return new ValueVector(bestSucc.getActions(), bestSucc.obj2, this, vf);
	}
	
//	public SingleStagePolicy bestAction = null; // prob over jTjA indexes (optPt returned by LP)
	
//	static class VVInfo {
//		public final NPoint optPoint;
//		public final BasicSolution bfs;
//		public VVInfo(NPoint optPoint, BasicSolution bfs) {
//			this.optPoint = optPoint;
//			this.bfs = bfs;
//		}
//		public VVInfo(double[] optPoint, BasicSolution bfs) {
//			this(new NPoint(optPoint), bfs);
//		}
//	}
//	public HashMap<Belief, VVInfo> lastInfo = new HashMap<Belief, VVInfo>();
	
	public Belief(double[] baseDist) {
		this.dist = new SparsePoint(baseDist);
		validateAndNormalize();
	}
	public Belief(SparsePoint baseDist) {
		this.dist = baseDist;
		validateAndNormalize();
	}
	@Override
	public Belief clone() {
		Belief b = new Belief(this.dist.clone());
		b.valueAtBelief = valueAtBelief;
		b.supportVector = supportVector;
		return b;
	}
	
	public static Belief makeUniformB(int dimensions) {
		return new Belief(new NPoint(dimensions, 1.0 / dimensions).values);
	}
	
	public double getProb(int stateIndex) {
		return dist.get(stateIndex);
	}
	public SparsePoint getProbDist() {
		return dist;
	}
	
	public double getValue() {
		return valueAtBelief;
	}
	public ValueVector getSupport() {
		return supportVector;
	}
	public void setSupport(ValueVector supportVector) {
		this.supportVector = supportVector;
	}
	
	private void validateAndNormalize() {
		// check to make sure this is a probability distribution (greater than 0, sums to one);
		double sum = 0;
		double sumNew = 0;
		Iterator<Entry<Integer, Double>> it = dist.iterator();
		while(it.hasNext()) {
			Entry<Integer, Double> e = it.next();
			Double value = e.getValue();
			assert(value >= -Globals.TOLERANCE && value <= 1.0 + Globals.TOLERANCE);
			if (value > 1.0 - Globals.TOLERANCE) {
				e.setValue(1.0);
				sumNew += 1.0;
			} else if (value < Globals.TOLERANCE) {
				it.remove();
			} else {
				sumNew += value;
			}
			sum += value;
		}
		assert(Math.abs(sum - 1.0) < Globals.TOLERANCE);
		dist.scale(1.0/sumNew);
	}
	
	public boolean validate(boolean[] validStateList) {
		assert(dist.isDistribution());
		for(Entry<Integer, Double> e : dist) {
			assert(validStateList[e.getKey()]);
		}
		return true;
	}

	public boolean updateIfBetter(ValueVector supportVector) {
		double possibleValue = supportVector.getValue(this);
		if(possibleValue > valueAtBelief) {		
			this.valueAtBelief = possibleValue;
			this.supportVector = supportVector;
			return true;
		}
		return false;
	}
	public void updateToBest(Collection<ValueVector> supportVectors) {
		this.valueAtBelief = Double.NEGATIVE_INFINITY;
		for(ValueVector v : supportVectors) {
			updateIfBetter(v);
		}
	}
	
	public Belief swapTypes(int[] typeSS, int agent, int t1, int t2, boolean[] validTypes) {
		SparsePoint newD = this.dist.clone();
		int[] numT = typeSS.clone();
		numT[agent] = 1;								      
		for(int[] jT : new CartesianIterator(numT)) {
			jT[agent] = t1;
			int jTi1 = JointIterator.getJointIndex(typeSS, jT);
			jT[agent] = t2;
			int jTi2 = JointIterator.getJointIndex(typeSS, jT);
			if(newD.get(jTi1) == 0.0 && newD.get(jTi2) == 0.0) {
				// nothing to swap, doesn't matter if they are invalid types
				continue;
			} else if((validTypes[jTi1] && validTypes[jTi2])) {
				double tmp = newD.get(jTi1);
				newD.set(jTi1, newD.get(jTi2));
				newD.set(jTi2, tmp);
			} else {
				// if we are trying to swap a valid type with an invalid type 
				// (and one has positive probability) then there is a problem
				return null;
			}
		}
		return new Belief(newD);
	}

	public static boolean softContains(Set<Belief> s, Belief b) {
		for(Belief b2 : s) {
			if(b.dist.softEquals(b2.dist, 0.0001))
				return true;
		}
		return false;
	}
	
//	public int findFirstNoneZero() {
//		for(int i=0; i<dimensions(); i++) {
//			if(values[i] != 0.0)
//				return i;
//		}
//		return -1;
//	}
	
	/**
	 * compute the type that has the closest euclidean distance
	 */
//	public Belief findClosest(Collection<Belief> bs) {
//		Belief best = null;
//		double bestSqDist = Double.MAX_VALUE;
//		for(Belief b : bs) {
//			double dist = this.sqDistance(b);
//			if(dist < bestSqDist) {
//				best = b;
//				bestSqDist = dist;
//			}
//		}
//		return best;
//	}
	
	public DecMDPSingleStagePolicy getCurPolicy() {
		return this.supportVector.nextActions;
	}
	
//	public String toNPointString() {
//		return super.toString();
//	}
	public String toSparseString() {
		return dist.toString();
	}
	
	@Override
	public String toString() {
		return "Belief:" + System.identityHashCode(this);
//		return getClass().getName() + "@" + Integer.toHexString(hashCode());
//		return "Belief at " +  "[value = " + valueAtBelief + "]" +
//			    ((values.length < 10)?super.toString():toSparseString()) ;
	}
	public int dimensions() {
		return dist.dimensions();
	}

}
